from _typeshed import Incomplete
from types import EllipsisType
from typing import Any, Generic, Self, SupportsIndex, overload, override
from typing_extensions import TypeVar, deprecated

import numpy as np
import numpy.typing as npt
from numpy._typing import (
    _AnyShape,
    _ArrayLike,
    _ArrayLikeBool_co,
    _ArrayLikeInt_co,
    _DTypeLike,
    _Shape,
)

###

_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], default=_AnyShape, covariant=True)
_DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype, default=np.dtype, covariant=True)

type _ArrayInt_co = npt.NDArray[np.integer | np.bool]

type _BoolContainer = container[Any, np.dtype[np.bool]]  # type: ignore[deprecated]
type _IntegralContainer = container[Any, np.dtype[np.bool | np.integer | np.object_]]  # type: ignore[deprecated]
type _RealContainer = container[Any, np.dtype[np.bool | np.integer | np.floating | np.timedelta64 | np.object_]]  # type: ignore[deprecated]
type _NumericContainer = container[Any, np.dtype[np.number | np.timedelta64 | np.object_]]  # type: ignore[deprecated]

type _ToIndexSlice = slice | EllipsisType | _ArrayInt_co | None
type _ToIndexSlices = _ToIndexSlice | tuple[_ToIndexSlice, ...]
type _ToIndex = SupportsIndex | _ToIndexSlice
type _ToIndices = _ToIndex | tuple[_ToIndex, ...]

###
# pyright: reportDeprecated = false

@deprecated("The numpy.lib.user_array.container class is deprecated and will be removed in a future version.")
class container(Generic[_ShapeT_co, _DTypeT_co]):
    array: np.ndarray[_ShapeT_co, _DTypeT_co]

    @overload
    def __init__(
        self,
        /,
        data: container[_ShapeT_co, _DTypeT_co] | np.ndarray[_ShapeT_co, _DTypeT_co],
        dtype: None = None,
        copy: bool = True,
    ) -> None: ...
    @overload
    def __init__[ScalarT: np.generic](
        self: container[Any, np.dtype[ScalarT]],
        /,
        data: _ArrayLike[ScalarT],
        dtype: None = None,
        copy: bool = True,
    ) -> None: ...
    @overload
    def __init__[ScalarT: np.generic](
        self: container[Any, np.dtype[ScalarT]],
        /,
        data: npt.ArrayLike,
        dtype: _DTypeLike[ScalarT],
        copy: bool = True,
    ) -> None: ...
    @overload
    def __init__(self, /, data: npt.ArrayLike, dtype: npt.DTypeLike | None = None, copy: bool = True) -> None: ...

    #
    def __complex__(self, /) -> complex: ...
    def __float__(self, /) -> float: ...
    def __int__(self, /) -> int: ...
    def __hex__(self, /) -> str: ...
    def __oct__(self, /) -> str: ...

    #
    @override
    def __eq__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
    @override
    def __ne__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]

    #
    def __lt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
    def __le__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
    def __gt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
    def __ge__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...

    #
    def __len__(self, /) -> int: ...

    # keep in sync with np.ndarray
    @overload
    def __getitem__(self, key: _ArrayInt_co | tuple[_ArrayInt_co, ...], /) -> container[_ShapeT_co, _DTypeT_co]: ...
    @overload
    def __getitem__(self, key: _ToIndexSlices, /) -> container[_AnyShape, _DTypeT_co]: ...
    @overload
    def __getitem__(self, key: _ToIndices, /) -> Any: ...
    @overload
    def __getitem__(self: container[Any, np.dtype[np.void]], key: list[str], /) -> container[_ShapeT_co, np.dtype[np.void]]: ...
    @overload
    def __getitem__(self: container[Any, np.dtype[np.void]], key: str, /) -> container[_ShapeT_co, np.dtype]: ...

    # keep in sync with np.ndarray
    @overload
    def __setitem__(self, index: _ToIndices, value: object, /) -> None: ...
    @overload
    def __setitem__(self: container[Any, np.dtype[np.void]], key: str | list[str], value: object, /) -> None: ...

    # keep in sync with np.ndarray
    @overload
    def __abs__[ShapeT: _Shape](
        self: container[ShapeT, np.dtype[np.complex64]], /
    ) -> container[ShapeT, np.dtype[np.float32]]: ...
    @overload
    def __abs__[ShapeT: _Shape](
        self: container[ShapeT, np.dtype[np.complex128]], /
    ) -> container[ShapeT, np.dtype[np.float64]]: ...
    @overload
    def __abs__[ShapeT: _Shape](
        self: container[ShapeT, np.dtype[np.complex192]], /
    ) -> container[ShapeT, np.dtype[np.float96]]: ...
    @overload
    def __abs__[ShapeT: _Shape](
        self: container[ShapeT, np.dtype[np.complex256]], /
    ) -> container[ShapeT, np.dtype[np.float128]]: ...
    @overload
    def __abs__[ContainerT: _RealContainer](self: ContainerT, /) -> ContainerT: ...

    #
    def __neg__[ContainerT: _NumericContainer](self: ContainerT, /) -> ContainerT: ...  # noqa: PYI019
    def __pos__[ContainerT: _NumericContainer](self: ContainerT, /) -> ContainerT: ...  # noqa: PYI019
    def __invert__[ContainerT: _IntegralContainer](self: ContainerT, /) -> ContainerT: ...  # noqa: PYI019

    # TODO(jorenham): complete these binary ops

    #
    def __add__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __radd__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __iadd__(self, other: npt.ArrayLike, /) -> Self: ...

    #
    def __sub__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __rsub__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __isub__(self, other: npt.ArrayLike, /) -> Self: ...

    #
    def __mul__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __rmul__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __imul__(self, other: npt.ArrayLike, /) -> Self: ...

    #
    def __mod__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __rmod__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __imod__(self, other: npt.ArrayLike, /) -> Self: ...

    #
    def __divmod__(self, other: npt.ArrayLike, /) -> tuple[Incomplete, Incomplete]: ...
    def __rdivmod__(self, other: npt.ArrayLike, /) -> tuple[Incomplete, Incomplete]: ...

    #
    def __pow__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __rpow__(self, other: npt.ArrayLike, /) -> Incomplete: ...
    def __ipow__(self, other: npt.ArrayLike, /) -> Self: ...

    #
    def __lshift__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.integer]]: ...
    def __rlshift__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.integer]]: ...
    def __ilshift__(self, other: _ArrayLikeInt_co, /) -> Self: ...

    #
    def __rshift__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.integer]]: ...
    def __rrshift__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.integer]]: ...
    def __irshift__(self, other: _ArrayLikeInt_co, /) -> Self: ...

    #
    @overload
    def __and__(self: _BoolContainer, other: _ArrayLikeBool_co, /) -> container[_AnyShape, np.dtype[np.bool]]: ...
    @overload
    def __and__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.bool | np.integer]]: ...
    __rand__ = __and__
    @overload
    def __iand__[ContainerT: _BoolContainer](self: ContainerT, other: _ArrayLikeBool_co, /) -> ContainerT: ...
    @overload
    def __iand__(self, other: _ArrayLikeInt_co, /) -> Self: ...

    #
    @overload
    def __xor__(self: _BoolContainer, other: _ArrayLikeBool_co, /) -> container[_AnyShape, np.dtype[np.bool]]: ...
    @overload
    def __xor__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.bool | np.integer]]: ...
    __rxor__ = __xor__
    @overload
    def __ixor__[ContainerT: _BoolContainer](self: ContainerT, other: _ArrayLikeBool_co, /) -> ContainerT: ...
    @overload
    def __ixor__(self, other: _ArrayLikeInt_co, /) -> Self: ...

    #
    @overload
    def __or__(self: _BoolContainer, other: _ArrayLikeBool_co, /) -> container[_AnyShape, np.dtype[np.bool]]: ...
    @overload
    def __or__(self, other: _ArrayLikeInt_co, /) -> container[_AnyShape, np.dtype[np.bool | np.integer]]: ...
    __ror__ = __or__
    @overload
    def __ior__[ContainerT: _BoolContainer](self: ContainerT, other: _ArrayLikeBool_co, /) -> ContainerT: ...
    @overload
    def __ior__(self, other: _ArrayLikeInt_co, /) -> Self: ...

    #
    @overload
    def __array__(self, /, t: None = None) -> np.ndarray[_ShapeT_co, _DTypeT_co]: ...
    @overload
    def __array__[DTypeT: np.dtype](self, /, t: DTypeT) -> np.ndarray[_ShapeT_co, DTypeT]: ...

    #
    @overload
    def __array_wrap__(self, arg0: npt.ArrayLike, /) -> container[_ShapeT_co, _DTypeT_co]: ...
    @overload
    def __array_wrap__[ShapeT: _Shape, DTypeT: np.dtype](
        self, a: np.ndarray[ShapeT, DTypeT], c: Any = ..., s: Any = ..., /
    ) -> container[ShapeT, DTypeT]: ...

    #
    def copy(self, /) -> Self: ...
    def tobytes(self, /) -> bytes: ...
    def byteswap(self, /) -> Self: ...
    def astype[ScalarT: np.generic](self, /, typecode: _DTypeLike[ScalarT]) -> container[_ShapeT_co, np.dtype[ScalarT]]: ...
